import numpy as np
import pandas as pd

# import matplotlib.pyplot as plt

from pgmpy.models import BayesianModel
from pgmpy.estimators import BayesianEstimator

"""
数据格式：
Pclass => 客舱等级(1/2/3等舱位)
Sex => 性别 male=1 female=0
Age => 年龄 插补后分0,1,2 代表 幼年（0-15） 成年（15-55） 老年（55-）
Fare => 船票价格 经聚类变0 1 2 代表少 多 很多
Cabin => 客舱号 清洗成有无此项，并发现有的生存率高
"""

# 读取训练集数据
train = pd.read_csv("./clean_train.csv")
# 读取测试集数据
test = pd.read_csv("./clean_test.csv")

# print(train.head())

# 人工构造贝叶斯网络结构
model = BayesianModel(
    [
        ("Age", "Survived"),
        ("Sex", "Survived"),
        ("Fare", "Pclass"),
        ("Pclass", "Survived"),
        ("Cabin", "Survived"),
    ]
)
# 使用训练集训练贝叶斯网络模型
model.fit(
    train, estimator=BayesianEstimator, prior_type="BDeu"
)  # default equivalent_sample_size=5

# for cpd in model.get_cpds():
#    print(cpd)

# 使用测试集进行预测
predict_data = test.drop(columns=["Survived"], axis=1)
y_pred = model.predict(predict_data)

print("实际结果:", y_pred["Survived"])
print("预测结果:", test["Survived"])
print("预测准确率:", (y_pred["Survived"] == test["Survived"]).sum() / len(test))

# 保存预测结果
y_pred.to_csv("test_predict.csv", index=0)
